#include "record.h"
#include <stdio.h>
#include <time.h>
#include "liteai/yolo.h"

Record::Record(std::string sn, int recsize)
{
	sn_ = sn;
	pkg_save_size_ = recsize;
	if (pkg_save_size_ == 0) {
		pkg_save_size_ = 240;
	}

	pkg_size_ = pkg_save_size_ * 3;

	fmt_ctx_ = NULL;
	
	pkt_list_ = (AVPacket**)malloc(sizeof(AVPacket*) * pkg_size_);
	memset(pkt_list_, 0, sizeof(AVPacket*) * pkg_size_);
	pkg_idx_ = 0;

	pkg_save_count_ = 0;
	
	saveing_ = false;

	for (int i = 0; i < pkg_size_; i++) {
		pkt_list_[i] = av_packet_alloc();
	}
}

Record::~Record()
{
	if (fmt_ctx_ != NULL) {
		avformat_close_input(&fmt_ctx_);
		fmt_ctx_ = NULL;
	}

	for (int i = 0; i < pkg_size_; i++) {
		av_packet_free(&pkt_list_[i]);
	}

	free(pkt_list_);
}

void Record::DetectObjsToSave(AVFormatContext* ictx, DetectObjsInfo* objs)
{
	if (saveing_) {
		return;
	}

	objsinof_ = objs;
	if (pkg_save_count_ > pkg_save_size_) {
		pkg_save_count_ = pkg_save_size_;
	}
	pkg_rec_idx_ = pkg_idx_;

	time_t now = time(NULL);
	struct tm* info = localtime(&now);

	char time_str[32];
	strftime(time_str, sizeof(time_str), "%Y%m%d%H%M%S", info);
	
	objs->recfile = std::string("./cache/") + sn_ + std::string("-") + time_str + std::string(".mp4");
	objs->imgfile = std::string("./cache/") + sn_ + std::string("-") + time_str + std::string(".jpg");
	
	strftime(time_str, 32, "%Y-%m-%d %H:%M:%S", info);
	objs->timestamp = time_str;

	filename_ = objs->recfile;

	if (!cv::imwrite(objs->imgfile.c_str(), objs->cvimg)) {
		fprintf(stderr, "img write to file error. %s\n", objs->imgfile.c_str());
		return;
	}

	if (fmt_ctx_ != NULL) {
		avformat_close_input(&fmt_ctx_);
	}

	avformat_alloc_output_context2(&fmt_ctx_, NULL, NULL, filename_.c_str());
	if (fmt_ctx_ == NULL) {
		fprintf(stderr, "could not create output context\n");
		return;
	}

	AVStream* stream = avformat_new_stream(fmt_ctx_, NULL);
	if (stream == NULL) {
		fprintf(stderr, "failed allocating output stream\n");
		return;
	}

	if (avcodec_parameters_copy(stream->codecpar, ictx->streams[0]->codecpar) < 0) {
		fprintf(stderr, "failed to copy codec parameters\n");
		return;
	}
	stream->codecpar->codec_tag = 0;

	time_base_ = ictx->streams[0]->time_base;

	av_dump_format(fmt_ctx_, 0, filename_.c_str(), 1);

	const AVOutputFormat* ofmt = fmt_ctx_->oformat;

	if (!(ofmt->flags & AVFMT_NOFILE)) {
		if (avio_open(&fmt_ctx_->pb, filename_.c_str(), AVIO_FLAG_WRITE) < 0) {
			fprintf(stderr, "could not open output file '%s'", filename_.c_str());
			return;
		}
	}

	if (avformat_write_header(fmt_ctx_, NULL) < 0) {
		fprintf(stderr, "error occurred when opening output file\n");
		return;
	}

	saveing_ = true;
	
}

AVPacket* Record::GetPkg(void) 
{
	this->SaveToFile();

	if (pkg_idx_ == pkg_size_) {
		pkg_idx_ = 0;
	}

	AVPacket* pkg = pkt_list_[pkg_idx_];
	
	pkg_idx_++;
	pkg_save_count_++;

	return pkg;
}

bool Record::IsSaveing(void) 
{
	return saveing_;
}

void Record::SaveToFile()
{
	if (!saveing_) {
		return;
	}
	
	if (pkg_save_count_ < pkg_save_size_*2) {
		return;
	}

	int start = pkg_rec_idx_ - pkg_save_size_;
	if (start < 0) {
		start = pkg_size_ + start;
	}

	bool key = false;

	int pts = 0;
	int dts = 0;

	for (int i = 0; i < pkg_save_size_ * 2; i++) {

		int ind = start + i;
		if (ind == pkg_size_) {
			start = -i;
			ind = 0;
		}
		AVPacket* pkt = pkt_list_[ind];
		if (pkt == NULL) {
			break;
		}

		if (!key) {
			if (!(pkt->flags & AV_PKT_FLAG_KEY)) {
				continue;
			}
			key = true;
			pts = pkt->pts;
			dts = pkt->dts;
		}

		pkt->pts -= pts;
		pkt->dts -= pts;

		//fprintf(stderr, "pts %d dts %d\n", pkt->pts, pkt->dts);

		if (av_interleaved_write_frame(fmt_ctx_, pkt) < 0) {
			fprintf(stderr, "error muxing packet %d\n", i);
		}
		//av_packet_unref(pkt);
	}

	av_write_trailer(fmt_ctx_);

	if (fmt_ctx_ != NULL) {
		avformat_close_input(&fmt_ctx_);
		fmt_ctx_ = NULL;
	}

	PostEvent();

	pkg_save_count_ = 0;
	pkg_rec_idx_ = 0;

	saveing_ = false;
}

void Record::PostEvent() 
{
	event_.Frie(objsinof_);
}
